﻿/*************************************************************************************
 * 工    具：  HK_MaterialTextureMatch
 * 
 * 描    述：  HK_MaterialTextureMatch 是一个 Unity 编辑器扩展工具，
 *            旨在简化材质球与其对应贴图之间的匹配过程。该工具支持多种渲染管线，包括内置渲染管线、URP 和 HDRP，并提供了灵活的贴图后缀匹配与自动分配功能。
 * 
 * 版    本：  V1.1
 * 作    者：  王凯、京产肠饭
 * 
 * 创    建：  2025/01/18  V1.0
 *
 * 更    新：  2025/02/06  V1.1
 *             1.功能优化：通过 EditorPrefs记录当前选择的渲染管线，避免每次打开窗口都要重新选择 
 *             2.功能优化：重新绘制 GUI 界面，提高用户体验
 * 
 * 链    接：  https://gitee.com/JingChanChangFan/hk_-unity-tools
*************************************************************************************/

using UnityEngine;
using UnityEditor;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System;
using Object = UnityEngine.Object;

namespace HKTools
{
    public class HK_MaterialTextureMatch : EditorWindow
    {
        enum RenderPipeline
        {
            BuildIn,
            URP,
            HDRP
        }

        /// <summary>
        /// 获取或设置当前渲染管线，通过EditorPrefs进行缓存。
        /// </summary>
        RenderPipeline CurrentPipeline
        {
            get
            {
                string pipelineString = EditorPrefs.GetString("CurrentPipeline", RenderPipeline.URP.ToString());
                return (RenderPipeline)Enum.Parse(typeof(RenderPipeline), pipelineString);
            }
            set
            {
                EditorPrefs.SetString("CurrentPipeline", value.ToString());
            }
        }

        List<Object> selectedFolders = new List<Object>();

        Dictionary<string, List<string>> textureSuffixes = new Dictionary<string, List<string>>
        {
            {"BaseMap", new List<string> {"_A", "_Albedo", "_BaseColor"}},
            {"MetallicMap", new List<string> {"_MS", "_Metallic", "_MetallicSmoothness","_Mask"}},
            {"AO", new List<string> {"_AO", "_Occlusion"}},
            {"Normal", new List<string> {"_N", "_Normal", "_NormalMap"}},
            {"Emission", new List<string> {"_E", "_Emission", "_Emissive"}}
        };

        Vector2 scrollPosition;
        bool isAutoEnableEmission = true;
        bool isReplaceM = false; // 新增的 toggle 选项
        bool isMergeBaseMapAndAO = false; // 新增选项

        [MenuItem("HKTools/Material Texture Match (材质球贴图匹配工具)", priority = HK_MenuPriority.HK_MaterialTextureMatch)]
        public static void ShowWindow()
        {
            GetWindow<HK_MaterialTextureMatch>("Material Texture Assigner").Show();
        }

        void OnGUI()
        {
            EditorGUILayout.Space();

            CurrentPipeline = (RenderPipeline)EditorGUILayout.EnumPopup("渲染管线：", CurrentPipeline);

            EditorGUILayout.Space();

            EditorGUILayout.LabelField("选择文件夹：", GUILayout.Width(100));

            scrollPosition = EditorGUILayout.BeginScrollView(scrollPosition, GUILayout.Height(selectedFolders.Count * 20 + 5));

            for (int i = 0; i < selectedFolders.Count; i++)
            {
                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.ObjectField(selectedFolders[i], typeof(Object), false);
                if (GUILayout.Button("X", GUILayout.Width(20)))
                {
                    selectedFolders.RemoveAt(i);
                    i--;
                }
                EditorGUILayout.EndHorizontal();
            }

            EditorGUILayout.EndScrollView();

            EditorGUILayout.Space();

            // 添加一个拖放区域
            Rect dropArea = GUILayoutUtility.GetRect(0.0f, 50.0f, GUILayout.ExpandWidth(true));
            GUI.Box(dropArea, "拖动文件夹到此处...");
            HandleFolderDraFgAndDrop(dropArea);

            if (GUILayout.Button("选择文件夹"))
            {
                HandleFolderSelect();
            }

            EditorGUILayout.Space();

            // 绘制一个分割线
            EditorGUI.DrawRect(EditorGUILayout.GetControlRect(false, 1), Color.gray);

            EditorGUILayout.Space();

            EditorGUILayout.LabelField("Texture 后缀：");

            foreach (var key in textureSuffixes.Keys.ToList())
            {
                EditorGUILayout.BeginVertical("box");
                EditorGUILayout.LabelField(key);
                string suffixes = string.Join(", ", textureSuffixes[key]);
                string newSuffixes = EditorGUILayout.TextField(suffixes);
                if (newSuffixes != suffixes)
                {
                    textureSuffixes[key] = new List<string>(newSuffixes.Split(',').Select(s => s.Trim()));
                }
                EditorGUILayout.EndVertical();
            }

            EditorGUILayout.Space();

            // 绘制一个分割线
            EditorGUI.DrawRect(EditorGUILayout.GetControlRect(false, 1), Color.gray);

            EditorGUILayout.Space();

            isAutoEnableEmission = EditorGUILayout.ToggleLeft("是否自动开启发光贴图", isAutoEnableEmission);
            isReplaceM = EditorGUILayout.ToggleLeft("是否将'_M'替换", isReplaceM);
            isMergeBaseMapAndAO = EditorGUILayout.ToggleLeft("是否合并 BaseMap 和 AO (HDRP)", isMergeBaseMapAndAO);

            EditorGUILayout.Space();

            if (GUILayout.Button("匹配贴图"))
            {
                HandleTexturesMatch();
            }
        }

        /// <summary>
        /// 处理拖放操作的方法。
        /// </summary>
        void HandleFolderDraFgAndDrop(Rect dropArea)
        {
            // 检查鼠标位置是否在拖放区域内，如果不在则直接返回。
            if (!dropArea.Contains(Event.current.mousePosition))
                return;

            // 检查当前事件类型是否为 DragUpdated。
            if (Event.current.type == EventType.DragUpdated)
            {
                // 设置拖放操作的视觉模式为复制。
                DragAndDrop.visualMode = DragAndDropVisualMode.Copy;
                // 标记事件已被使用，防止其他处理程序处理该事件。
                Event.current.Use();
            }

            // 检查当前事件类型是否为 DragPerform。
            if (Event.current.type == EventType.DragPerform)
            {
                // 设置拖放操作的视觉模式为复制。
                DragAndDrop.visualMode = DragAndDropVisualMode.Copy;
                // 接受拖放操作。
                DragAndDrop.AcceptDrag();

                // 遍历拖放的对象引用。
                foreach (Object draggedObject in DragAndDrop.objectReferences)
                {
                    // 获取拖放对象的路径。
                    string path = AssetDatabase.GetAssetPath(draggedObject);
                    // 检查路径是否为目录。
                    if (Directory.Exists(path))
                    {
                        // 将目录添加到选中的文件夹列表中。
                        if (selectedFolders.Contains(draggedObject))
                            continue;

                        selectedFolders.Add(draggedObject);
                    }
                }

                // 标记事件已被使用，防止其他处理程序处理该事件。
                Event.current.Use();
            }
        }

        /// <summary>
        /// 处理文件夹选择
        /// </summary>
        void HandleFolderSelect()
        {
            string[] folderPaths = EditorUtility.OpenFolderPanel("Select Folders", "Assets", "").Split('\n');

            foreach (string folderPath in folderPaths)
            {
                if (string.IsNullOrEmpty(folderPath))
                    continue;

                string relativePath = "Assets" + folderPath.Substring(Application.dataPath.Length);
                Object folder = AssetDatabase.LoadAssetAtPath<Object>(relativePath);
                if (folder != null && !selectedFolders.Contains(folder))
                {
                    selectedFolders.Add(folder);
                }
            }
        }

        /// <summary>
        /// 处理材质贴图匹配
        /// </summary>
        void HandleTexturesMatch()
        {
            foreach (Object folder in selectedFolders)
            {
                string folderPath = AssetDatabase.GetAssetPath(folder);
                string[] materialGuids = AssetDatabase.FindAssets("t:Material", new[] { folderPath });

                foreach (string materialGuid in materialGuids)
                {
                    string materialPath = AssetDatabase.GUIDToAssetPath(materialGuid);
                    Material material = AssetDatabase.LoadAssetAtPath<Material>(materialPath);

                    if (material == null)
                        continue;

                    string materialName = Path.GetFileNameWithoutExtension(materialPath);
                    AssignMaterial(material, materialName, folderPath);
                }
            }

            AssetDatabase.SaveAssets();
            AssetDatabase.Refresh();
        }

        void AssignMaterial(Material material, string materialName, string folderPath)
        {
            foreach (var textureSuffix in textureSuffixes)
            {
                string baseTextureName = isReplaceM ? materialName.Replace("_M", "") : materialName;
                Texture foundTexture = null;

                foreach (string suffix in textureSuffix.Value)
                {
                    string textureName = baseTextureName + suffix;
                    string[] textureGuids = AssetDatabase.FindAssets(textureName + " t:Texture", new[] { folderPath });

                    if (textureGuids.Length > 0)
                    {
                        string texturePath = AssetDatabase.GUIDToAssetPath(textureGuids[0]);
                        foundTexture = AssetDatabase.LoadAssetAtPath<Texture>(texturePath);
                        break;
                    }
                }
                // 如果严格匹配未找到，开始进行模糊匹配
                if (foundTexture == null)
                {
                    foreach (string suffix in textureSuffix.Value)
                    {
                        string textureName = baseTextureName + suffix;
                        string candidateTextureName = FindClosestMatchingTextureName(baseTextureName, textureSuffix.Value, folderPath);

                        HKDebug.Log($"模糊搜索 :<Color=Red> {candidateTextureName} </Color>");

                        if (!string.IsNullOrEmpty(candidateTextureName))
                        {
                            string[] textureGuids = AssetDatabase.FindAssets(candidateTextureName + " t:Texture", new[] { folderPath });

                            if (textureGuids.Length > 0)
                            {
                                string texturePath = AssetDatabase.GUIDToAssetPath(textureGuids[0]);
                                foundTexture = AssetDatabase.LoadAssetAtPath<Texture>(texturePath);
                                HKDebug.Log($"模糊匹配到材质 '{materialName}' 的 {textureSuffix.Key} 纹理: {texturePath}");
                                break;
                            }
                        }
                    }
                }

                if (foundTexture != null)
                {
                    switch (CurrentPipeline)
                    {
                        case RenderPipeline.BuildIn:
                            AssignTextureBuiltIn(material, textureSuffix.Key, foundTexture);
                            break;
                        case RenderPipeline.URP:
                            AssignTextureURP(material, textureSuffix.Key, foundTexture);
                            break;
                        case RenderPipeline.HDRP:
                            AssignTextureHDRP(material, textureSuffix.Key, foundTexture);
                            break;
                    }
                }
                else
                {
                    HKDebug.LogWarning($"在文件夹 '{folderPath}' 中未找到材质 '{materialName}' 的类型为 '{textureSuffix.Key}' 的纹理。");
                }
            }
        }

        string FindClosestMatchingTextureName(string baseTextureName, List<string> suffixes, string folderPath)
        {
            string[] supportedExtensions = new string[] { "*.png", "*.jpg", "*.tga", "*.jpeg", "*.tif", "*.tiff" };
            List<string> allTextureFiles = new List<string>();

            // 收集所有支持格式的纹理文件
            foreach (string extension in supportedExtensions)
            {
                allTextureFiles.AddRange(Directory.GetFiles(folderPath, extension, SearchOption.AllDirectories));
            }

            string closestMatch = null;

            foreach (string filePath in allTextureFiles)
            {
                string fileName = Path.GetFileNameWithoutExtension(filePath).ToLower();
                string baseTextureLower = baseTextureName.ToLower();

                foreach (string suffix in suffixes)
                {
                    string expectedName = (baseTextureLower + suffix).ToLower();

                    // 精确匹配
                    if (fileName.Equals(expectedName))
                    {
                        return fileName; // 找到精确匹配，立即返回
                    }

                    // 对于长度大于4的后缀，允许一个字母的差异或少一个字母
                    if (suffix.Length > 4)
                    {
                        if (IsCloseMatch(fileName, expectedName))
                        {
                            closestMatch = fileName;

                            break; // 找到接近匹配，继续检查其他文件
                        }
                    }
                }

                if (closestMatch != null)
                {
                    break; // 如果找到接近匹配，停止搜索
                }
            }

            return closestMatch;
        }

        bool IsCloseMatch(string actual, string expected)
        {
            // 将两个字符串都转换为小写
            actual = actual.ToLower();
            expected = expected.ToLower();

            // 检查长度差异
            int lengthDiff = Math.Abs(actual.Length - expected.Length);
            if (lengthDiff > 1)
                return false;

            if (lengthDiff == 0)
            {
                // 如果长度相同，直接计算不同字符的数量
                return actual.Zip(expected, (a, e) => a == e).Count(eq => !eq) <= 1;
            }
            else
            {
                // 如果长度相差1，使用编辑距离算法
                return ComputeEditDistance(actual, expected) <= 1;
            }
        }

        int ComputeEditDistance(string s1, string s2)
        {
            int[,] dp = new int[s1.Length + 1, s2.Length + 1];

            for (int i = 0; i <= s1.Length; i++)
                dp[i, 0] = i;

            for (int j = 0; j <= s2.Length; j++)
                dp[0, j] = j;

            for (int i = 1; i <= s1.Length; i++)
            {
                for (int j = 1; j <= s2.Length; j++)
                {
                    if (s1[i - 1] == s2[j - 1])
                        dp[i, j] = dp[i - 1, j - 1];
                    else
                        dp[i, j] = Math.Min(dp[i - 1, j - 1], Math.Min(dp[i - 1, j], dp[i, j - 1])) + 1;
                }
            }

            return dp[s1.Length, s2.Length];
        }

        void AssignTextureBuiltIn(Material material, string textureType, Texture texture)
        {
            switch (textureType)
            {
                case "BaseMap":
                    material.SetTexture("_MainTex", texture);
                    break;
                case "MetallicMap":
                    material.SetTexture("_MetallicGlossMap", texture);
                    break;
                case "AO":
                    material.SetTexture("_OcclusionMap", texture);
                    break;
                case "Normal":
                    material.SetTexture("_BumpMap", texture);
                    material.EnableKeyword("_NORMALMAP");
                    break;
                case "Emission":
                    material.SetTexture("_EmissionMap", texture);
                    if (isAutoEnableEmission)
                    {
                        material.EnableKeyword("_EMISSION");
                        material.globalIlluminationFlags = MaterialGlobalIlluminationFlags.RealtimeEmissive;
                    }
                    break;
            }
        }

        void AssignTextureURP(Material material, string textureType, Texture texture)
        {
            switch (textureType)
            {
                case "BaseMap":
                    material.SetTexture("_BaseMap", texture);
                    break;
                case "MetallicMap":
                    material.SetTexture("_MetallicGlossMap", texture);
                    break;
                case "AO":
                    material.SetTexture("_OcclusionMap", texture);
                    break;
                case "Normal":
                    material.SetTexture("_BumpMap", texture);
                    material.EnableKeyword("_NORMALMAP");
                    break;
                case "Emission":
                    material.SetTexture("_EmissionMap", texture);
                    if (isAutoEnableEmission)
                    {
                        material.EnableKeyword("_EMISSION");
                        material.globalIlluminationFlags = MaterialGlobalIlluminationFlags.RealtimeEmissive;
                    }
                    break;
            }
        }

        void AssignTextureHDRP(Material material, string textureType, Texture texture)
        {
            switch (textureType)
            {
                case "BaseMap":
                    if (isMergeBaseMapAndAO)
                    {
                        // 查找 AO 纹理
                        Texture aoTexture = FindTextureByType(material.name, "AO", AssetDatabase.GetAssetPath(material));
                        if (aoTexture != null)
                        {
                            // 合并 BaseMap 和 AO
                            Texture2D mergedTexture = MergeBaseMapAndAO((Texture2D)texture, (Texture2D)aoTexture);
                            material.SetTexture("_BaseColorMap", mergedTexture);
                            HKDebug.Log($"为材质 '{material.name}' 合并了 BaseMap 和 AO");
                        }
                        else
                        {
                            material.SetTexture("_BaseColorMap", texture);
                            HKDebug.LogWarning($"未找到材质 '{material.name}' 的 AO 纹理。使用未合并的 BaseMap。");
                        }
                    }
                    else
                    {
                        material.SetTexture("_BaseColorMap", texture);
                    }
                    break;
                case "MetallicMap":
                    material.SetTexture("_MaskMap", texture);
                    break;
                case "AO":
                    if (!isMergeBaseMapAndAO)
                    {
                        // 如果不合并，AO 可能会被用作 MaskMap 的一部分
                        HKDebug.LogWarning($"在HDRP中，AO通常是MaskMap的一部分。请考虑将其与其他纹理合并。");
                    }
                    break;
                case "Normal":
                    material.SetTexture("_NormalMap", texture);
                    material.EnableKeyword("_NORMALMAP");
                    break;
                case "Emission":
                    material.SetTexture("_EmissiveColorMap", texture);
                    if (isAutoEnableEmission)
                    {
                        material.EnableKeyword("_EMISSION");
                        material.globalIlluminationFlags = MaterialGlobalIlluminationFlags.RealtimeEmissive;
                        material.SetFloat("_UseEmissiveIntensity", 1);
                        material.SetColor("_EmissiveColor", Color.white);
                        material.SetFloat("_EmissiveIntensityUnit", 0);
                        material.SetFloat("_EmissiveExposureWeight", 1);
                    }
                    break;
            }
        }

        Texture FindTextureByType(string materialName, string textureType, string folderPath)
        {
            string baseTextureName = isReplaceM ? materialName.Replace("_M", "") : materialName;
            foreach (string suffix in textureSuffixes[textureType])
            {
                string textureName = baseTextureName + suffix;
                string[] textureGuids = AssetDatabase.FindAssets(textureName + " t:Texture", new[] { folderPath });
                if (textureGuids.Length > 0)
                {
                    string texturePath = AssetDatabase.GUIDToAssetPath(textureGuids[0]);
                    return AssetDatabase.LoadAssetAtPath<Texture>(texturePath);
                }
            }
            return null;
        }

        Texture2D MergeBaseMapAndAO(Texture2D baseMap, Texture2D aoMap)
        {
            if (baseMap == null || aoMap == null)
            {
                HKDebug.LogError("BaseMap 或 AO Map 为空！");
                return null;
            }

            int width = baseMap.width;
            int height = baseMap.height;

            // 确保两个纹理大小一致
            if (aoMap.width != width || aoMap.height != height)
            {
                HKDebug.LogWarning("AO Map的尺寸与Base Map的尺寸不匹配，正在调整AO Map的尺寸。");
                aoMap = ResizeTexture(aoMap, width, height);
            }

            Texture2D mergedTexture = new Texture2D(width, height, TextureFormat.RGBA32, false);

            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    Color baseColor = baseMap.GetPixel(x, y);
                    Color aoColor = aoMap.GetPixel(x, y);

                    // 将 AO 信息存储到 Alpha 通道
                    Color mergedColor = new Color(baseColor.r, baseColor.g, baseColor.b, aoColor.grayscale);
                    mergedTexture.SetPixel(x, y, mergedColor);
                }
            }

            mergedTexture.Apply();
            return mergedTexture;
        }

        Texture2D ResizeTexture(Texture2D source, int newWidth, int newHeight)
        {
            RenderTexture rt = RenderTexture.GetTemporary(newWidth, newHeight);
            Graphics.Blit(source, rt);
            RenderTexture previous = RenderTexture.active;
            RenderTexture.active = rt;

            Texture2D resizedTexture = new Texture2D(newWidth, newHeight);
            resizedTexture.ReadPixels(new Rect(0, 0, newWidth, newHeight), 0, 0);
            resizedTexture.Apply();

            RenderTexture.active = previous;
            RenderTexture.ReleaseTemporary(rt);

            return resizedTexture;
        }
    }
}